package com.lou1052.redis.jedis.ext.client.cluster;

import com.lou1052.redis.jedis.ext.client.single.JedisExt;
import com.lou1052.redis.jedis.ext.client.single.JedisExtPool;
import org.apache.commons.pool2.PooledObject;
import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
import redis.clients.jedis.HostAndPort;
import redis.clients.jedis.exceptions.JedisConnectionException;

import java.util.*;
import java.util.function.Supplier;

public class JedisExtSlotBasedConnectionHandler extends JedisClusterExtConnectionHandler {

  public JedisExtSlotBasedConnectionHandler(Set<HostAndPort> nodes,
                                            final GenericObjectPoolConfig poolConfig, int timeout, Supplier<Deque<PooledObject<JedisExt>>> dequeSupplier) {
    this(nodes, poolConfig, timeout, timeout,dequeSupplier);
  }

  public JedisExtSlotBasedConnectionHandler(Set<HostAndPort> nodes,
                                            final GenericObjectPoolConfig poolConfig, int connectionTimeout, int soTimeout, Supplier<Deque<PooledObject<JedisExt>>> dequeSupplier) {
    super(nodes, poolConfig, connectionTimeout, soTimeout,dequeSupplier);
  }

  public JedisExt getConnection() {
    // In antirez's redis-rb-cluster implementation,
    // getRandomConnection always return valid connection (able to
    // ping-pong)
    // or exception if all connections are invalid

    List<JedisExtPool> pools = getShuffledNodesPool();

    for (JedisExtPool pool : pools) {
      JedisExt jedis = null;
      try {
        jedis = pool.getResource();

        if (jedis == null) {
          continue;
        }

        String result = jedis.ping();

        if (result.equalsIgnoreCase("pong")) return jedis;

        pool.returnBrokenResource(jedis);
      } catch (JedisConnectionException ex) {
        if (jedis != null) {
          pool.returnBrokenResource(jedis);
        }
      }
    }

    throw new JedisConnectionException("no reachable node in cluster");
  }

  @Override
  public JedisExt getConnectionFromSlot(int slot) {
    JedisExtPool connectionPool = cache.getSlotPool(slot);
    if (connectionPool != null) {
      // It can't guaranteed to get valid connection because of node
      // assignment
      return connectionPool.getResource();
    } else {
      return getConnection();
    }
  }

  private List<JedisExtPool> getShuffledNodesPool() {
    List<JedisExtPool> pools = new ArrayList<>();
    pools.addAll(cache.getNodes().values());
    Collections.shuffle(pools);
    return pools;
  }

}
